from typing import Any, Dict, List, Optional, Tuple, Type
from psycopg2.extensions import connection
from psycopg2.extras import RealDictCursor


class DBUtils:
    @staticmethod
    def query_one(conn: connection, sql: str, params: Tuple[Any, ...], dto_class: Optional[Type] = None) -> Optional[Any]:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(sql, params)
            row = cur.fetchone()
            if row and dto_class:
                return dto_class(**row)
            return row

    @staticmethod
    def query_all(conn: connection, sql: str, params: Tuple[Any, ...] = (), dto_class: Optional[Type] = None) -> List[Any]:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(sql, params)
            rows = cur.fetchall()
            if dto_class:
                return [dto_class(**row) for row in rows]
            return rows

    @staticmethod
    def insert(conn: connection, sql: str, params: Tuple[Any, ...]) -> int:
        with conn.cursor() as cur:
            cur.execute(sql, params)
            conn.commit()
            return cur.rowcount

    @staticmethod
    def insert_and_return_id(conn: connection, sql: str, params: Tuple[Any, ...], id_column: str = "id") -> Any:
        with conn.cursor() as cur:
            cur.execute(f"{sql} RETURNING {id_column}", params)
            new_id = cur.fetchone()[0]
            conn.commit()
            return new_id

    @staticmethod
    def update(conn: connection, sql: str, params: Tuple[Any, ...]) -> int:
        with conn.cursor() as cur:
            cur.execute(sql, params)
            conn.commit()
            return cur.rowcount

    @staticmethod
    def delete(conn: connection, sql: str, params: Tuple[Any, ...]) -> int:
        with conn.cursor() as cur:
            cur.execute(sql, params)
            conn.commit()
            return cur.rowcount

    @staticmethod
    def paginate_query(
        conn: connection,
        base_sql: str,
        count_sql: str,
        conditions: List[str],
        params: List[Any],
        page: int,
        size: int,
        dto_class: Optional[Type] = None,
        order_by: str = "id DESC"
    ) -> Tuple[int, List[Any]]:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            where_sql = ""
            if conditions:
                where_sql = " WHERE " + " AND ".join(conditions)

            # 查询总数
            full_count_sql = count_sql + where_sql
            cur.execute(full_count_sql, params)
            total = cur.fetchone()["count"]

            # 查询分页数据
            full_query_sql = base_sql + where_sql + f" ORDER BY {order_by} LIMIT %s OFFSET %s"
            page_params = params + [size, (page - 1) * size]
            cur.execute(full_query_sql, page_params)

            rows = cur.fetchall()
            if dto_class:
                return total, [dto_class(**row) for row in rows]
            return total, rows